//      build_fom.cc
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#include <boost/program_options.hpp>
#include <boost/date_time/posix_time/posix_time.hpp> 
#include <boost/foreach.hpp>

#include "helper/stl_helper.h"
#include "helper/ublas_helper.h"
#include "helper/ml_data.h"
#include "helper/rand_helper.h"
#include "helper/dna_processing.h"
#include "helper/suffix_array.h"

#include "cds_vom.h"

extern boost::program_options::variables_map vm;

boost::posix_time::ptime now() ;
boost::posix_time::time_duration duration() ;

void compute_likelihood(Sequences& S, const std::string& llout, VOM& model)
{
  FileAccessWrite(out, llout);
  int j = 1;
  BOOST_FOREACH(const std::string& s, S) {
    double ll = model.log_prob(s);
    out << ">" << j << "\n" << ll << ENDL;
    j++;
  }
}

void compute_likelihood_lambda(Sequences& S, const std::string& llout, VOM& model)
{
  FileAccessWrite(out, llout);
  int j = 1;
  BOOST_FOREACH(const std::string& s, S) {
    double ll = model.log_prob_lambda(s);
    out << ">" << j << "\n" << ll << ENDL;
    j++;
  }
}

void build_fom()
{
  std::string inFile = VM(vm, "input", std::string);
  vi_t orders = VM(vm, "order", std::vector<int>);
  int period = VM(vm, "period", int);
  MSG("++ Build FOM input=" << inFile << " period=" << period);

  // read sequences and sort suffixes
  Sequences S;
  S.fasta_read(inFile);
  MSG("read " << sz(S) << " sequences.");
  sastring_vec_t SA(sz(S));
  std::copy(all_element(S), SA.begin());
  MSG("suffix array of " << sz(SA) << " sequences built.");

  // likelihood files
  std::vector<Sequences> llSeqs;
  std::vector<std::string> ll, llout;
  if (vm.count("ll") && vm.count("llout")) {
    ll = VM(vm, "ll", std::vector<std::string>);
    llout = VM(vm, "llout", std::vector<std::string>);
    if (sz(ll) != sz(llout)) 
      FATAL_ERROR("Check --ll and --llout, should have same number of file names");
    for (int i = 0; i < sz(ll); i++) {
      Sequences seqs;
      seqs.fasta_read(ll.at(i));
      llSeqs.pb(seqs);
      MSG("Read likelihood sequences ll=" << ll.at(i) << " llout=" << llout.at(i));
    }
  }

  // build models
  BOOST_FOREACH(int order, orders) {
    VOM model("ACGT", period);
    MSG("order=" << order << " duration = " << duration() << " now = " << now());
    //model.build_fom(SA, order);
    model.build_fom_flat(SA, order);
    MSG("built " << " duration = " << duration() << " now = " << now());

    if (vm.count("output")) {
      std::string outFile = VM(vm, "output", std::string) + "." + tos(order);
      FileAccessWrite(out, outFile);
      model.print_node(out);
    }
    
    if (vm.count("fom_stats")) {
      std::string statsFile = VM(vm, "fom_stats", std::string) + "." + tos(order);
      model.depth();
      model.stats_print(statsFile);
    }

    if (vm.count("fom_ll")  && sz(llSeqs) > 0) { 
      MSG("FOM likelihood");
      for (int i = 0; i < sz(llSeqs); i++) {
	FileAccessWrite(out, llout.at(i) + "." + tos(order));
	int j = 1;
	BOOST_FOREACH(const std::string& s, llSeqs[i]) {
	  double ll = model.log_prob(s);
	  out << ">" << j << "\n" << ll << ENDL;
	  j++;
	}
      }
    }

    if (vm.count("imm_xi_ll") && sz(llSeqs) > 0) {
      MSG("IMM likelihood");

      model.T_ = vm.count("imm_T") ? VM(vm, "imm_T", int) : 100;
      model.compute_lambda_xi_square();
      for (int i = 0; i < sz(llSeqs); i++) {
	FileAccessWrite(out, llout.at(i) + ".xi." + tos(order));
	int j = 1;
	BOOST_FOREACH(const std::string& s, llSeqs[i]) {
	  double ll = model.log_prob_lambda(s);
	  out << ">" << j << "\n" << ll << ENDL;
	  j++;
	}
      }
    } // imm-xi

    if (vm.count("imm_di_ll") && sz(llSeqs) > 0) {
      MSG("IMM likelihood (deleted interpolation)");
      VOM di_model("ACGT", period);
      di_model.copy(model);
      di_model.compute_lambda_deleted_interpolation();
      for (int i = 0; i < sz(llSeqs); i++) {
	FileAccessWrite(out, llout.at(i) + ".di." + tos(order));
	int j = 1;
	BOOST_FOREACH(const std::string& s, llSeqs[i]) {
	  double ll = di_model.log_prob_lambda(s);
	  out << ">" << j << "\n" << ll << ENDL;
	  j++;
	}
      }
    } // imm-xi

    if (vm.count("prune_ll") && vm.count("eta")) {
      std::vector<double> etas = VM(vm, "eta", std::vector<double>);
      BOOST_FOREACH(double eta, etas) {
	MSG("Pruned likelihood eta="<<eta);
	VOM pruned_model("ACGT", period);
	pruned_model.copy(model);
	pruned_model.prune(SA, eta);

	if (vm.count("fom_stats")) {
	  std::string statsFile = VM(vm, "fom_stats", std::string) 
	    + ".prune." + tos(eta) + "." + tos(order);
	  pruned_model.depth();
	  pruned_model.stats_print(statsFile);
	}
	for (int i = 0; i < sz(llSeqs); i++)
	  compute_likelihood(llSeqs[i], llout.at(i) + ".prune." + tos(eta) + "." + tos(order), 
			     pruned_model);
      }
    }

    if (vm.count("prune_xi_ll") && vm.count("eta")) {
      std::vector<double> etas = VM(vm, "eta", std::vector<double>);
      model.T_ = vm.count("imm_T") ? VM(vm, "imm_T", int) : 100;
      model.compute_lambda_xi_square();
      BOOST_FOREACH(double eta, etas) {
	MSG("Pruned IMM likelihood eta="<<eta);
	VOM pruned_model("ACGT", period);
	pruned_model.copy(model);
	pruned_model.prune_lambda(SA, eta);

	if (vm.count("fom_stats")) {
	  std::string statsFile = VM(vm, "fom_stats", std::string) 
	    + ".prune.xi." + tos(eta) + "." + tos(order);
	  pruned_model.depth();
	  pruned_model.stats_print(statsFile);
	}
	for (int i = 0; i < sz(llSeqs); i++)
	  compute_likelihood_lambda(llSeqs[i], llout.at(i) + ".prune.xi." + tos(eta) + "." + tos(order), 
				    pruned_model);
      }
    }

  } // BOOST_FOREACH(order)
}
